package com.wistech.search.solr.extend;

import java.io.IOException;
import java.util.Map;
import java.util.Set;

import org.apache.lucene.document.Document;
import org.apache.lucene.queryparser.classic.ParseException;
import org.apache.lucene.search.Query;
import org.apache.solr.common.SolrException;
import org.apache.solr.common.SolrInputField;
import org.apache.solr.common.util.NamedList;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.response.SolrQueryResponse;
import org.apache.solr.search.DocIterator;
import org.apache.solr.search.DocSet;
import org.apache.solr.search.QParser;
import org.apache.solr.search.SolrIndexSearcher;
import org.apache.solr.update.AddUpdateCommand;
import org.apache.solr.update.CommitUpdateCommand;
import org.apache.solr.update.DeleteUpdateCommand;
import org.apache.solr.update.MergeIndexesCommand;
import org.apache.solr.update.RollbackUpdateCommand;
import org.apache.solr.update.UpdateHandler;
import org.apache.solr.update.processor.UpdateRequestProcessor;
import org.apache.solr.update.processor.UpdateRequestProcessorFactory;

import com.google.common.collect.Sets;
import com.wistech.commons.utils.spring.SpringContextHolder;

/**
 * 自定义的RunUpdateProcessorFactory,用于处理文档更新时的外部数据及score计算
 * 
 */
public class WisRunUpdateProcessorFactory extends UpdateRequestProcessorFactory {
	private static final String DOC_HANDLER = "doc-handler";
	private static final String EXTERNAL_FIELD_PREFIX = "prefix";
	private String prefix = "live_";
	private DocumentHandlerManager documentHandlerManager;

	/*
	 * (non-Javadoc)
	 * 
	 * @see
	 * org.apache.solr.update.processor.UpdateRequestProcessorFactory#init(org
	 * .apache.solr.common.util.NamedList)
	 */
	@SuppressWarnings("rawtypes")
	@Override
	public void init(NamedList args) {
		super.init(args);

		// 获取外部字段的前缀
		String prefix = (String) args.get(EXTERNAL_FIELD_PREFIX);
		if (prefix != null) {
			this.prefix = prefix;
		}
		// 获取外部数据处理者配置项
		String documentHandlerManagerName = (String) args.get(DOC_HANDLER);
		if (documentHandlerManagerName != null) {
			documentHandlerManagerName = documentHandlerManagerName.replaceAll(
					"\t|\r|\n", "");// 去除字符串中的回车/换行符/制表符
			try {
				documentHandlerManager = (DocumentHandlerManager) SpringContextHolder
						.getBean(documentHandlerManagerName);
			} catch (Exception e) {
				throw new SolrException(SolrException.ErrorCode.SERVER_ERROR,
						e.getMessage());

			}
		}
	}

	@Override
	public UpdateRequestProcessor getInstance(SolrQueryRequest req,
			SolrQueryResponse rsp, UpdateRequestProcessor next) {
		return new WisRunUpdateProcessor(req, next, this);
	}

	/**
	 * @return the prefix
	 */
	protected String getPrefix() {
		return prefix;
	}

	/**
	 * @return the documentHandlerManager
	 */
	protected DocumentHandlerManager getDocumentHandlerManager() {
		return documentHandlerManager;
	}

}

class WisRunUpdateProcessor extends UpdateRequestProcessor {
	private final SolrQueryRequest req;
	private final UpdateHandler updateHandler;
	private DocumentHandlerManager documentHandlerManager;
	private String idFieldName;
	private String prefix;

	public WisRunUpdateProcessor(SolrQueryRequest req,
			UpdateRequestProcessor next, WisRunUpdateProcessorFactory factory) {
		super(next);
		this.req = req;
		this.updateHandler = req.getCore().getUpdateHandler();
		idFieldName = req.getSchema().getUniqueKeyField().getName();
		this.documentHandlerManager = factory.getDocumentHandlerManager();
		this.prefix = factory.getPrefix();
	}

	@Override
	public void processAdd(AddUpdateCommand cmd) throws IOException {
		// 添加索引
		if (needIndexing(cmd)) {
			updateHandler.addDoc(cmd);
		}
		// 保存外部数据
		if (documentHandlerManager != null) {
			SolrIndexContext context = new SolrIndexContext(this.req);
			context.setDocument(cmd.getSolrInputDocument());
			context.setRequest(this.req);
			context.handleDocumentWithPrefix(prefix);
			documentHandlerManager.add(context);
		}
		super.processAdd(cmd);
	}

	private boolean needIndexing(AddUpdateCommand cmd) {
		String idFieldName = this.req.getSchema().getUniqueKeyField().getName();
		for (Map.Entry<String, SolrInputField> entry : cmd
				.getSolrInputDocument().entrySet()) {
			String fieldName = entry.getKey();
			if (!fieldName.equals(idFieldName) && !fieldName.startsWith(prefix)) {
				return true;
			}
		}
		return false;
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see
	 * org.apache.solr.update.processor.UpdateRequestProcessor#processDelete
	 * (org.apache.solr.update.DeleteUpdateCommand)
	 */
	@Override
	public void processDelete(DeleteUpdateCommand cmd) throws IOException {
		// 删除外部数据
		if (cmd.id != null) {
			documentHandlerManager.delete(cmd.id);
		} else {
			Set<String> fields = Sets.newHashSetWithExpectedSize(1);
			fields.add(this.idFieldName);
			QParser parser;
			try {
				parser = QParser.getParser(cmd.query, null, req);
				Query q = parser.getQuery();
				SolrIndexSearcher searcher = req.getSearcher();
				DocSet docSet = req.getSearcher().getDocSet(q);
				DocIterator iter = docSet.iterator();
				while (iter.hasNext()) {
					int docid = iter.nextDoc();
					Document doc = searcher.doc(docid, fields);
					String id = req.getSchema().getUniqueKeyField(doc)
							.stringValue();
					if (id != null) {
						documentHandlerManager.delete(id);
					}

				}
			} catch (ParseException e) {
				e.printStackTrace();
			}

		}
		// 删除索引
		if (cmd.id != null) {
			updateHandler.delete(cmd);
		} else {
			updateHandler.deleteByQuery(cmd);
		}
		super.processDelete(cmd);
	}

	@Override
	public void processMergeIndexes(MergeIndexesCommand cmd) throws IOException {
		updateHandler.mergeIndexes(cmd);
		super.processMergeIndexes(cmd);
	}

	@Override
	public void processCommit(CommitUpdateCommand cmd) throws IOException {
		updateHandler.commit(cmd);
		super.processCommit(cmd);
	}

	/**
	 * @since Solr 1.4
	 */
	@Override
	public void processRollback(RollbackUpdateCommand cmd) throws IOException {
		updateHandler.rollback(cmd);
		super.processRollback(cmd);
	}
}
